package org.moonlight.redis.limit.interceptor;

import lombok.extern.slf4j.Slf4j;
import org.moonlight.redis.limit.annotation.AccessLimit;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.OutputStream;
import java.lang.reflect.Method;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.TimeUnit;

/**
 * 访问拦截器
 * @author Moonlight
 */
@Slf4j
@Component
public class AccessLimitInterceptor implements HandlerInterceptor {

    private final RedisTemplate<String, Object> redisTemplate;

    @Autowired
    public AccessLimitInterceptor(RedisTemplate<String, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        try {
            if (handler instanceof HandlerMethod) {
                HandlerMethod handlerMethod = (HandlerMethod) handler;
                // 获取待执行的方法, 对于加了 AccessLimit 注解的方法，进行限流校验
                Method method = handlerMethod.getMethod();
                AccessLimit accessLimit = method.getAnnotation(AccessLimit.class);
                if (accessLimit != null) {
                    // 获取注解配置参数
                    int maxCount = accessLimit.maxCount();
                    int second = accessLimit.second();
                    AccessLimit.AccessLimitStrategy strategy = accessLimit.strategy();
                    // 默认以 ContextPath:ServletPath 拼接为Key
                    String key = request.getContextPath() + ":" + request.getServletPath();

                    if (AccessLimit.AccessLimitStrategy.ACCESS_IP.equals(strategy)) {
                        // 如果配置是以访问者IP进行限制, 则使用 访问者IP:ContextPath:ServletPath 拼接为Key
                        key = request.getRemoteAddr() + ":" + request.getContextPath() + ":" + request.getServletPath();
                    }
                    // 获取访问次数
                    Integer accessCnt = (Integer) redisTemplate.opsForValue().get(key);
                    // 第一次访问或者之前的限流已经过期了
                    if (accessCnt == null) {
                        redisTemplate.opsForValue().set(key, 1, second, TimeUnit.SECONDS);
                        return true;
                    }
                    if (accessCnt < maxCount) {
                        redisTemplate.opsForValue().increment(key);
                    } else {
                        // 超过最大次数
                        log.warn("[AccessLimitInterceptor preHandle] 服务繁忙, 请稍后重试. maxCount[{}] second[{}] strategy[{}] key[{}] accessCnt[{}]"
                                , maxCount, second, strategy, key, accessCnt);
                        respServiceBusy(response);
                        return false;
                    }
                }
            }
        } catch (Exception e) {
            log.error("[AccessLimitInterceptor preHandle] 进行访问限流处理时出错. ", e);
        }
        return true;
    }

    private void respServiceBusy(HttpServletResponse response) throws IOException {
        response.setContentType("application/json;charset=UTF-8");
        try (OutputStream out = response.getOutputStream()) {
            out.write("服务繁忙, 请稍后重试".getBytes(StandardCharsets.UTF_8));
        }
    }
}